#ifndef EXP__HUFFMAN_TREE_HPP
#define EXP__HUFFMAN_TREE_HPP

#include <iostream>
#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>

namespace my_data
{

const int N = 1010;

struct Node {
char ch;  //节点对应的字符
int weight; // 结点权值
int left; // 左儿子编号
int right; // 右儿子编号
int parent; // 父节点编号
bool isLeaf; // 是否为叶子结点
};

class HuffmanCode {
public:
/**
 * @brief 构造函数
 */ 
  explicit HuffmanCode(std::unordered_map<char,int> freq_table);
/**
 * @brief 输入需要编码的字符串
 */ 
  void input(std::string & string);
/**
 * @brief 生成哈夫曼树
 */   
  void initHuffmanTree();
/**
 * @brief 编码
 */ 
  void enCode();
/**
 * @brief 输出编码的结果
 */ 
  std::string output();
private:

private:
  std::vector<Node> huffman_tree_; // 哈夫曼树
  int n_; // 结点数量
  std::string str_; // 输入的字符串
  std::vector<int> freq_; // 字符出现频率
  std::vector<std::string> code_; //编码
};

HuffmanCode::HuffmanCode(std::unordered_map<char,int> freq_table)
{
  n_ = 0;
  huffman_tree_.reserve(2 * N);
  freq_.reserve(N);
  code_.reserve(N);
  for(int i=0; i<N; ++i)
  {
    freq_[i] = 0;
  }
  for(auto pair : freq_table)
  {
    freq_[(int)pair.first] = pair.second; 
  }
}

void HuffmanCode::input(std::string &string)
{
  str_ = string;
}

void HuffmanCode::initHuffmanTree()
{
  // 初始化每个结点
  for (int i = 0; i < N; i++) 
  {
    if (freq_[i] != 0) 
    {
      n_++;
      huffman_tree_[n_].weight = freq_[i];
      huffman_tree_[n_].isLeaf = true;
      huffman_tree_[n_].ch = (char)i;
    }
  }
  // 构建哈夫曼树
  for (int i = n_ + 1; i <= 2 * n_; i++) 
  {
    // 找权重最小的数和权重第二小的树
    int min1 = 0, min2 = 0;
    for (int j = 1; j < i; j++) 
    {
      if (huffman_tree_[j].parent == 0) 
      {
        if (min1 == 0 || huffman_tree_[j].weight < huffman_tree_[min1].weight) 
        {
          min2 = min1;
          min1 = j;
        }
        else if (min2 == 0 || huffman_tree_[j].weight < huffman_tree_[min2].weight) 
        {
          min2 = j;
        }
      }
    }
    // 两棵树合成一颗树
    huffman_tree_[i].weight = huffman_tree_[min1].weight + huffman_tree_[min2].weight;
    huffman_tree_[i].left = min1;
    huffman_tree_[i].right = min2;
    huffman_tree_[min1].parent = i;
    huffman_tree_[min2].parent = i;
    huffman_tree_[i].isLeaf = false;
  }
}

void HuffmanCode::enCode()
{
  for (int i = 1; i <= n_; i++) 
  {
    std::string s;
    int cur = i;
    int parent = huffman_tree_[i].parent;
    // 从当前结点往上追溯，碰到根节点停止
    while (parent != 0) 
    {
      if (huffman_tree_[parent].left == cur) 
      {
        s += "0"; // 左儿子则记为0
      }
      else 
      {
        s += "1"; // 右儿子则记为1
      }
      cur = parent;
      parent = huffman_tree_[parent].parent;
    }
    std::reverse(s.begin(), s.end()); // 反转字符串
    code_[i] = s;
  }
}

std::string HuffmanCode::output()
{
  std::string s = "|";
  for(char c : str_)
  {
    for(int i=1; i<=n_; ++i)
    {
      if(c == huffman_tree_[i].ch)
      {
        s += code_[i];
        s += '|';
      }
    }
  }
  for (int i = 1; i <= n_; i++) {
    std::cout << huffman_tree_[i].ch << ": " << code_[i] << std::endl;
  }
  std::cout<<"\""<<str_<<"\""<<" 的编码为:\n"<<s<<std::endl;
  return s;
}

} //namespace my_data

#endif  //EXP__HUFFMAN_TREE_HPP